#include <assert.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#include "../src/multimodel.c"

static bool
fuzzy_equal (float a,
             float b)
{
  return (fabs (a - b) <= 1e-5 * (fabs (a) + fabs (b)));
}

static bool
fuzzy_equal_array (const float* a,
                   const float* b,
                   size_t n)
{
  for (unsigned int i = 0; i < n; ++i)
    {
      if (! fuzzy_equal (a[i], b[i]))
        {
          return false;
        }
    }

  return true;
}

static void
test_model_once ()
{
  const char* pathname = tmpnam (NULL);
  unsigned int n_items = 1.0 + 50.0 * drand48 ();
  unsigned int n_raw_labels = 1.0 + 10.0 * drand48 ();
  unsigned int max_raw_labels = n_raw_labels * drand48 ();
  unsigned int n_workers = 1.0 + 50.0 * drand48 ();
  unsigned int rank = 1.0 + 10.0 * drand48 ();
  unsigned int effective_n_labels = 0;

  for (unsigned int m = 0; m <= max_raw_labels; ++m)
    {
      effective_n_labels += binomial (n_raw_labels, m);
    }

  float logpriorz[n_raw_labels];
  float loga[n_workers * rank];
  float gamma[rank];
  float c[effective_n_labels * rank];
  float d[effective_n_labels * rank];

  for (unsigned int i = 0; i < n_raw_labels; ++i)
    {
      logpriorz[i] = drand48 ();
    }

  float effective_logpriorz[effective_n_labels];

  for (unsigned int i = 0; i < effective_n_labels; ++i)
    {
      unsigned int labels[n_raw_labels];
      unsigned int depth = index_to_set (i, labels, n_raw_labels);
      float logpz = 0.0;

      for (unsigned int j = 0; j < depth; ++j)
        {
          logpz += logpriorz[labels[j]];
        }

      effective_logpriorz[i] = logpz;
    }

  safe_exponentiate_normalize (effective_logpriorz, effective_n_labels);

  for (unsigned int i = 0; i < effective_n_labels; ++i)
    {
      effective_logpriorz[i] = log (effective_logpriorz[i]);
    }

  memset (loga, 0, sizeof (loga));
  memset (gamma, 0, sizeof (gamma));
  memset (c, 0, sizeof (c));
  memset (d, 0, sizeof (d));

  MultiModel* model = multi_model_creat (pathname,
                                         n_items,
                                         n_raw_labels,
                                         max_raw_labels,
                                         n_workers,
                                         rank,
                                         logpriorz);

  assert (model->n_items == n_items);
  assert (model->n_labels == effective_n_labels);
  assert (model->n_raw_labels == n_raw_labels);
  assert (model->max_raw_labels == max_raw_labels);
  assert (model->n_workers == n_workers);
  assert (model->rank == rank);
  assert (fuzzy_equal_array (model->logpriorz, 
                             effective_logpriorz,
                             effective_n_labels));
  assert (fuzzy_equal_array (model->rawlogpriorz, 
                             logpriorz,
                             n_raw_labels));
  assert (memcmp (model->loga,
                  loga,
                  n_workers * rank * sizeof (float)) == 0);
  assert (memcmp (model->gamma,
                  gamma,
                  rank * sizeof (float)) == 0);
  assert (memcmp (model->c, c, effective_n_labels * rank * sizeof (float)) != 0);
  assert (memcmp (model->d, d, effective_n_labels * rank * sizeof (float)) != 0);

  for (unsigned int i = 0; i < n_workers * rank; ++i)
    {
      model->loga[i] = loga[i] = drand48 ();
    }

  for (unsigned int i = 0; i < rank; ++i)
    {
      model->gamma[i] = gamma[i] = drand48 ();
    }

  for (unsigned int i = 0; i < effective_n_labels * rank; ++i)
    {
      model->c[i] = c[i] = drand48 ();
      model->d[i] = d[i] = drand48 ();
    }

  multi_model_close (model);

  model = multi_model_open (pathname, true);

  assert (model->n_items == n_items);
  assert (model->n_labels == effective_n_labels);
  assert (model->n_raw_labels == n_raw_labels);
  assert (model->max_raw_labels == max_raw_labels);
  assert (model->n_workers == n_workers);
  assert (model->rank == rank);
  assert (fuzzy_equal_array (model->logpriorz, 
                             effective_logpriorz,
                             effective_n_labels));
  assert (fuzzy_equal_array (model->rawlogpriorz, 
                             logpriorz,
                             n_raw_labels));
  assert (memcmp (model->loga,
                  loga,
                  n_workers * rank * sizeof (float)) == 0);
  assert (memcmp (model->gamma,
                  gamma,
                  rank * sizeof (float)) == 0);
  assert (fuzzy_equal_array (model->c, c, effective_n_labels * rank));
  assert (fuzzy_equal_array (model->d, d, effective_n_labels * rank));

  multi_model_close (model);
  unlink (pathname);
}

static void
test_model (void)
{
  for (unsigned int i = 0; i < 100; ++i)
    {
      test_model_once ();
    }
}

int 
main (void) 
{
  srand48 (69);

  test_model ();

  return 0;
}
